import random
from copy import deepcopy
from typing import Set, List

import networkx as nx

from model.action import *
from model.state import State
from model.transition import Transition


class CPTEM(nx.MultiDiGraph):
    # current_state:Optional[State] = None
    # firing:Optional[Transition] = None
    def __init__(self, entity: Optional[str] = None, attr: Optional[dict] = None):
        super().__init__()
        self.entity = entity
        self.current_state = None
        self.firing = None
        self.attr = attr


    def cptem_entity_type(self):
        if self.entity is None:
            return None
        if isinstance(self.entity, str):
            if self.entity.startswith("n"):
                return "node"
            if self.entity.startswith("s"):
                return "service"
            if self.entity.startswith("a"):
                return "application"
        if isinstance(self.entity, tuple):
            if self.entity[0].startswith("n"):
                return "physical edge"
            if self.entity[0].startswith("s"):
                return "service edge"
        raise TypeError("entity must be str or tuple")

    def add_state(self, state:State):
        self.add_node(state)

    def add_states(self, states:Set[State]):
        for state in states:
            self.add_state(state)
    
    def add_transition_rule(self, rule:Transition):
        self.add_edge(rule.source_state,
            rule.target_state,
            label = rule)
    
    def get_state_by_name(self, name:str) -> State:
        for state in self.states():
            if state.name == name:
                return state
            
    def get_transition_rule_by_name(self, name:str) -> Transition:
        for rule in self.get_transition_rules():
            if rule.name == name:
                return rule
    
    def get_transition_rule_by_states(self, source_state:State, target_state:State) -> List[Transition]:
        l = self.get_edge_data(source_state, target_state)
        return [v["label"] for _, v in l.items()]


    def add_transition_rules(self, rules:Set[Transition]):
        for rule in rules:
            self.add_transition_rule(rule)
    
    def get_transition_rules(self) -> Set[Transition]:
        return set(self.edges(data=True))
    
    def initial_current_state(self) :
        if len(self.nodes) == 0:
            self.current_state = None
        else:
            for state in self.states():
                if state.is_initial:
                    self.current_state = state
                    break
    
    def get_current_state(self) -> State:
        return self.current_state
    
    def set_current_state(self, state:State):
        if state in self.states():
            self.current_state = state
        else:
            raise ValueError("invalid state")

    def get_current_attr(self) -> dict:
        return self.attr

    def get_firing_transition(self) -> Transition:
        return self.firing
    
    def states(self) -> Set[State]:
        s = set()
        for v in self.nodes():
            s.add(v)
        return s
    
    def find_enable_transitions(self, act: Optional[Action]) -> List[Transition]:
        if self.get_current_state() is None:
            print("no current state")
            return []

        if self.firing is not None:
            return []

        enabled_ts = []
        out_states = list(self.successors(self.get_current_state()))
        for out_state in out_states:
            ts =  self.get_transition_rule_by_states(self.get_current_state(), out_state)
            for t in ts:
                if t.probability < 0 or t.probability > 1:
                    raise ValueError("probability must be at [0, 1]")
                if t.probability == 0:
                    continue
                if t.probability <= 1:
                    temp = random.random()
                    if temp > t.probability:
                        continue
                    else:
                        if t.guard is None:
                            enabled_ts.append(t)
                        else:
                            if act is not None:
                                if t.guard(act, self.attr):
                                    enabled_ts.append(t)
        return enabled_ts
    
    def firing_transition(self, transition: Transition):
        if self.firing is not None:
            raise ValueError("already firing")
        else:
            self.firing = transition

    def inhibit_transition(self):
        if self.firing is None:
            raise ValueError("there is no firing transition")
        else:
            self.firing = None

    def fired_transition_update_cptem(self, act: Action):
        if self.firing is None:
            raise ValueError("there is no firing transition")
        else:
            firing_transition = self.get_firing_transition()
            if self.get_current_state() != firing_transition.source_state:
                raise ValueError("firing transition is not from current state")
            else:
                self.firing = None
                self.set_current_state(firing_transition.get_target_state())
                if firing_transition.update is not None:
                    self.attr = firing_transition.update(act, self.attr)


    # def draw_straight_line(self):
    #
    #     figsize = (8, 6)
    #
    #     # create a new figure
    #     fig = plt.figure(figsize=figsize)
    #
    #     # draw the graph
    #     pos = nx.spring_layout(self)
    #     node_colors = ["lightblue" if n == self.get_current_state() else "lightgrey" for n in self.nodes()]
    #     nx.draw_networkx(self, pos, with_labels=True, node_size=5000, node_color=node_colors)
    #
    #     # Draw the edge labels
    #     edge_labels = {(u, v): str(d['label']) for u, v, d in self.edges(data=True)}
    #     nx.draw_networkx_edge_labels(self, pos, edge_labels=edge_labels)
    #
    #     # Show the plot
    #     plt.show()
    #
    # def draw(self):
    #
    #
    #     # Create a new figure
    #     fig = plt.figure()
    #
    #     # Draw the graph
    #     pos = nx.spring_layout(self)
    #     node_colors = ["lightblue" if n == self.get_current_state() else "lightgrey" for n in self.nodes()]
    #     nx.draw_networkx_nodes(self, pos, node_size=5000, node_color=node_colors)
    #
    #     for i, (u, v, d) in enumerate(self.edges(data=True)):
    #         label = str(d['label'])
    #         nx.draw_networkx_edges(self, pos,label=label, arrows=True, edgelist=[(u, v)], edge_color='gray', connectionstyle=f'arc3,rad=0.2')
    #         nx.draw_networkx_edge_labels(self, pos, edge_labels={(u, v): label}, label_pos=0.5, font_size=5, font_color='black')
    #
    #     # Show the plot
    #     plt.show()
    #
    
    def to_dot(self, name: str = "cptem.dot"):
        import graphviz
        # import pydot

        cptem = deepcopy(self)
        current_state = cptem.get_current_state()

        # Convert the graph to a DOT file
        graph_dot = nx.drawing.nx_pydot.to_pydot(self)
        graph_dot.get_node(str(current_state))[0].set_color("red")
        gv = graphviz.Source(graph_dot.to_string())

        # Display the Graphviz object
        gv.render(name, format='png')
    
